package avlTree;

public class AVLTree {
    public static class TreeNode {
        public int val;
        public int bf; // 平衡因子

        public TreeNode parent; // 父节点
        public TreeNode left; // 左子树
        public TreeNode right;// 右子树

        public TreeNode(int val) {
            this.val = val;
        }
    }

    public TreeNode root;

    public boolean insert(int val) {
        TreeNode node = new TreeNode(val);
        if(root == null) {
            root = node;
            return true;
        }

        // 添加
        TreeNode parent = null, cur = root;
        while(cur != null) {
            parent = cur;
            if(cur.val < val) {
                cur = cur.right;
            } else if(cur.val == val) {
                return false;
            } else {
                cur = cur.left;
            }
        }
        // cur = null
        if(parent.val < val) {
            parent.right = node;
        } else {
            parent.left = node;
        }
        node.parent = parent;

        // 更新平衡因子
        cur = node;
        while(parent != null) {
            // 先更新一下 parent 的平衡因子
            if(parent.left == cur) {
                parent.bf--;
            } else {
                parent.bf++;
            }

            // 判断平衡因子
            if(parent.bf == 0) {
                // 已经平衡
                break;
            } else if(parent.bf == 1 || parent.bf == -1) {
                // 向上传
                cur = parent;
                parent = cur.parent;
            } else {
                if(parent.bf == 2) {
                    // 右树太高
                    if(cur.bf == 1) {
                        rotateLeft(parent);
                    } else {
                        rotateRL(parent);
                        // cur.bf = -1
                    }
                } else {
                    // parent.bf = -2 左树太高
                    if(cur.bf == 1) {
                        rotateLR(parent);
                    } else {
                        // cur.bf = -1 左旋
                        rotateRight(parent);
                    }
                }
                break;
            }
        }
        return true;
    }

    /**
     * 左右双旋
     * @param parent
     */
    private void rotateLR(TreeNode parent) {
        TreeNode subL = parent.left, subLR = subL.right;
        int bf = subLR.bf;

        rotateLeft(subL);
        rotateRight(parent);

        subLR.bf = 0;
        if(bf == -1) {
            subL.bf = 0;
            parent.bf = 1;
        } else {
            subL.bf = -1;
            parent.bf = 0;
        }
    }


    /**
     * 右左双旋
     * @param parent
     */
    private void rotateRL(TreeNode parent) {
        TreeNode subR = parent.right, subRL = subR.left;
        int bf = subRL.bf;

        rotateRight(subR);
        rotateLeft(parent);


        subRL.bf = 0;
        if(bf == -1) {
            subR.bf = -1;
            parent.bf = 0;
        } else {
            subR.bf = 1;
            parent.bf = 0;
        }
    }

    /**
     * 左单旋
     * @param parent
     */
    private void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right, subRL = subR.left, pParent = parent.parent;

        parent.right = subRL;
        parent.parent = subR;

        if(subRL != null) {
            subRL.parent = parent;
        }

        subR.left = parent;
        subR.parent = pParent;

        subR.bf = 0;
        parent.bf = 0;

        if(parent == root) {
            root = subR;
            return;
        }
        if(pParent.left == parent) {
            pParent.left = subR;
        } else {
            pParent.right = subR;
        }
    }

    /**
     * 右单旋
     * @param parent
     */
    private void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left, subLR = subL.right;

        subL.parent = parent.parent;
        subL.right = parent.right;
        if(subLR != null) {
            // 没有右节点时
            subLR.parent = parent;
        }
        parent.left = subLR;
        TreeNode pParent = parent.parent;
        parent.parent = subL;

        // 将平衡因子置为 0
        parent.bf = 0;
        subL.bf = 0;

        // 判断 root 是不是 parent
        if(root == parent) {
            root = subL;
            return;
        }
        // 更新 pParent 的子树（如果有）
        if(pParent.left == parent) {
            pParent.left = subL;
        } else {
            pParent.right = subL;
        }
    }

}
